#include "ContentBrowserExtensions.h"
#include "Modules/ModuleManager.h"
#include "Misc/PackageName.h"
#include "Textures/SlateIcon.h"
#include "Framework/Commands/UIAction.h"
#include "Framework/MultiBox/MultiBoxExtender.h"
#include "Framework/MultiBox/MultiBoxBuilder.h"
#include "EditorStyleSet.h"
#include "Engine/Texture2D.h"
#include "AssetToolsModule.h"
#include "IContentBrowserSingleton.h"
#include "ContentBrowserModule.h"
#include "PixelShader.h"
#include "RenderObject.h"
#include "RenderTexture2D.h"
#include "VertexShader.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "Kismet2/KismetEditorUtilities.h"

#define LOCTEXT_NAMESPACE "SoftRendererEditor"

static FContentBrowserMenuExtender_SelectedAssets ContentBrowserExtenderDelegate;
static FDelegateHandle ContentBrowserExtenderDelegateHandle;

//////////////////////////////////////////////////////////////////////////
// FContentBrowserSelectedAssetExtensionBase

struct FContentBrowserSelectedAssetExtensionBase
{
public:
	TArray<struct FAssetData> SelectedAssets;

public:
	virtual void Execute() {}
	virtual ~FContentBrowserSelectedAssetExtensionBase() {}
};

//////////////////////////////////////////////////////////////////////////
// FCreateRenderObjectFromStaticMeshExtension

struct FCreateRenderObjectFromStaticMeshExtension : public FContentBrowserSelectedAssetExtensionBase
{
	FCreateRenderObjectFromStaticMeshExtension()
	{
	}
	
	virtual void Execute() override
	{
		const FAssetToolsModule& AssetToolsModule = FModuleManager::Get().LoadModuleChecked<FAssetToolsModule>("AssetTools");
		
		for (auto AssetIt = SelectedAssets.CreateConstIterator(); AssetIt; ++AssetIt)
		{
			const FAssetData& AssetData = *AssetIt;
			if (const UStaticMesh* StaticMesh = Cast<UStaticMesh>(AssetData.GetAsset()))
			{
				// Create the sprite
				FString Name;
				FString PackageName;
				
				const FString DefaultSuffix = TEXT("_RenderObject");
				AssetToolsModule.Get().CreateUniqueAssetName(StaticMesh->GetOutermost()->GetName(), DefaultSuffix, /*out*/ PackageName, /*out*/ Name);

				UPackage* Package = CreatePackage( *PackageName);
				if (ensure(Package))
				{
					UBlueprint* NewBP = CastChecked<UBlueprint>(FKismetEditorUtilities::CreateBlueprint(URenderObject::StaticClass(), Package, FName(Name), EBlueprintType::BPTYPE_Normal, 
						UBlueprint::StaticClass(), UBlueprintGeneratedClass::StaticClass(), NAME_None));

					if (const TSubclassOf<UObject> GeneratedClass = NewBP->GeneratedClass)
					{
						if (URenderObject* RenderObject = GeneratedClass->GetDefaultObject<URenderObject>())
						{
							const FStaticMeshLODResources& StaticMeshLOD = StaticMesh->RenderData->LODResources[0];
							
							const int32 NumWedges = StaticMeshLOD.IndexBuffer.GetNumIndices();
							const int32 NumVertexPositions = StaticMeshLOD.VertexBuffers.PositionVertexBuffer.GetNumVertices();

							//Create the vertex
							for (int32 VertexIndex = 0; VertexIndex < NumVertexPositions; ++VertexIndex)
							{
								FRenderObjectVertex Vertex;
								Vertex.Position = StaticMeshLOD.VertexBuffers.PositionVertexBuffer.VertexPosition(VertexIndex);

								if (VertexIndex < static_cast<int32>(StaticMeshLOD.VertexBuffers.ColorVertexBuffer.GetNumVertices()))
								{
									Vertex.VertexColor = StaticMeshLOD.VertexBuffers.ColorVertexBuffer.VertexColor(VertexIndex).ReinterpretAsLinear();
								}

								for (int32 UVIndex = 0, UVCount = StaticMeshLOD.VertexBuffers.StaticMeshVertexBuffer.GetNumTexCoords(); UVIndex < UVCount; ++UVIndex)
								{
									FVector2D UV = StaticMeshLOD.VertexBuffers.StaticMeshVertexBuffer.GetVertexUV(VertexIndex, UVIndex);
									Vertex.UVs.Emplace(UV);
								}

								Vertex.Normal = StaticMeshLOD.VertexBuffers.StaticMeshVertexBuffer.VertexTangentZ(VertexIndex);
								Vertex.Tangent = StaticMeshLOD.VertexBuffers.StaticMeshVertexBuffer.VertexTangentX(VertexIndex);
								
								RenderObject->Vertices.Emplace(Vertex);
							}

							for (int32 Index = 0; Index < NumWedges; ++Index)
							{
								RenderObject->Indices.Add(StaticMeshLOD.IndexBuffer.GetIndex(Index));
							}

							RenderObject->Material.VertexShaderClass = UVertexShader::StaticClass();
							RenderObject->Material.PixelShaderClass = UPixelShader::StaticClass();
						}
					}

					FBlueprintEditorUtils::MarkBlueprintAsModified(NewBP);
					
					// Need to make sure we compile with the new source code
					// FKismetEditorUtilities::CompileBlueprint(NewBP);
				}

			}
		}
 
	}
};

//////////////////////////////////////////////////////////////////////////
// FCreateRenderTexture2DFromTextureExtension

struct FCreateRenderTexture2DFromTextureExtension : public FContentBrowserSelectedAssetExtensionBase
{
	FCreateRenderTexture2DFromTextureExtension()
	{
	}
	
	virtual void Execute() override
	{
		const FAssetToolsModule& AssetToolsModule = FModuleManager::Get().LoadModuleChecked<FAssetToolsModule>("AssetTools");
		
		for (auto AssetIt = SelectedAssets.CreateConstIterator(); AssetIt; ++AssetIt)
		{
			const FAssetData& AssetData = *AssetIt;
			if (UTexture2D* Texture2D = Cast<UTexture2D>(AssetData.GetAsset()))
			{
				const ETextureSourceFormat TextureFormat = Texture2D->Source.GetFormat();
				if (TextureFormat != ETextureSourceFormat::TSF_BGRA8)
				{
					UE_LOG(LogTemp, Error, TEXT("不支持的贴图格式,只支持BGRA8的数据格式"));
					continue;
				}
				
				// Create the sprite
				FString Name;
				FString PackageName;
				
				const FString DefaultSuffix = TEXT("_RenderTexture2D");
				AssetToolsModule.Get().CreateUniqueAssetName(Texture2D->GetOutermost()->GetName(), DefaultSuffix, /*out*/ PackageName, /*out*/ Name);

				UPackage* Package = CreatePackage( *PackageName);
				if (ensure(Package))
				{
					UBlueprint* NewBP = CastChecked<UBlueprint>(FKismetEditorUtilities::CreateBlueprint(URenderTexture2D::StaticClass(), Package, FName(Name), EBlueprintType::BPTYPE_Normal, 
						UBlueprint::StaticClass(), UBlueprintGeneratedClass::StaticClass(), NAME_None));

					if (const TSubclassOf<UObject> GeneratedClass = NewBP->GeneratedClass)
					{
						if (URenderTexture2D* RenderTexture2D = GeneratedClass->GetDefaultObject<URenderTexture2D>())
						{
							if (Texture2D->AddressX == TextureAddress::TA_Clamp)
							{
								RenderTexture2D->AddressX = ERenderTextureAddressMode::TA_Clamp;
							}
							else if (Texture2D->AddressX == TextureAddress::TA_Mirror)
							{
								RenderTexture2D->AddressX = ERenderTextureAddressMode::TA_Mirror;
							}
							else if (Texture2D->AddressX == TextureAddress::TA_Wrap)
							{
								RenderTexture2D->AddressX = ERenderTextureAddressMode::TA_Wrap;
							}
						
							if (Texture2D->AddressY == TextureAddress::TA_Clamp)
							{
								RenderTexture2D->AddressY = ERenderTextureAddressMode::TA_Clamp;
							}
							else if (Texture2D->AddressY == TextureAddress::TA_Mirror)
							{
								RenderTexture2D->AddressY = ERenderTextureAddressMode::TA_Mirror;
							}
							else if (Texture2D->AddressY == TextureAddress::TA_Wrap)
							{
								RenderTexture2D->AddressY = ERenderTextureAddressMode::TA_Wrap;
							}

							if (Texture2D->Filter == TextureFilter::TF_Bilinear)
							{
								RenderTexture2D->Filter = ERenderTextureFilter::TF_Bilinear;
							}
							else if (Texture2D->Filter == TextureFilter::TF_Nearest)
							{
								RenderTexture2D->Filter = ERenderTextureFilter::TF_Nearest;
							}
							else if (Texture2D->Filter == TextureFilter::TF_Trilinear)
							{
								RenderTexture2D->Filter = ERenderTextureFilter::TF_Trilinear;
							}

							RenderTexture2D->bSRGB = Texture2D->SRGB;

							TArray64<uint8> RawData;
							Texture2D->Source.GetMipData(RawData, 0);

							FRenderTexture2DMipMap MipData;
							MipData.SizeX = Texture2D->Source.GetSizeX();
							MipData.SizeY = Texture2D->Source.GetSizeY();

							MipData.TextureData.SetNumZeroed(MipData.SizeX * MipData.SizeY * Texture2D->Source.GetBytesPerPixel(TextureFormat));
							FMemory::Memcpy(MipData.TextureData.GetData(), RawData.GetData(), MipData.SizeX * MipData.SizeY * Texture2D->Source.GetBytesPerPixel(TextureFormat));
							
							RenderTexture2D->Mips.Emplace(MipData);
						}
					}

					FBlueprintEditorUtils::MarkBlueprintAsModified(NewBP);
					
					// Need to make sure we compile with the new source code
					// FKismetEditorUtilities::CompileBlueprint(NewBP);
				}
			}
		}
	}
};

//////////////////////////////////////////////////////////////////////////
// FSoftRendererContentBrowserExtensions_Impl

class FSoftRendererContentBrowserExtensions_Impl
{
public:
	static void ExecuteSelectedContentFunctor(TSharedPtr<FContentBrowserSelectedAssetExtensionBase> SelectedAssetFunctor)
	{
		SelectedAssetFunctor->Execute();
	}

	static void CreateSoftRendererActionsSubMenu(FMenuBuilder& MenuBuilder, TArray<FAssetData> SelectedAssets, bool bAnyStaticMeshes, bool bAnyTextures)
	{
		if (bAnyStaticMeshes)
		{
			MenuBuilder.AddSubMenu(
				LOCTEXT("SoftRendererActionsSubMenuLabel", "SoftRenderer Actions"),
				LOCTEXT("SoftRendererActionsSubMenuToolTip", "SoftRenderer-related actions for static meshes."),
				FNewMenuDelegate::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::PopulateSoftRendererActionsMenu, SelectedAssets),
				false
			);
		}

		if (bAnyTextures)
		{
			MenuBuilder.AddSubMenu(
				LOCTEXT("SoftRendererActionsSubMenuLabel", "SoftRenderer Actions"),
				LOCTEXT("SoftRendererActionsSubMenuToolTip", "SoftRenderer-related actions for textures."),
				FNewMenuDelegate::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::PopulateRenderTexture2DActionsMenu, SelectedAssets),
				false
			);
		}
	}

	static void PopulateSoftRendererActionsMenu(FMenuBuilder& MenuBuilder, TArray<FAssetData> SelectedAssets)
	{
		const TSharedPtr<FCreateRenderObjectFromStaticMeshExtension> RenderObjectCreatorFunctor = MakeShareable(new FCreateRenderObjectFromStaticMeshExtension());
		RenderObjectCreatorFunctor->SelectedAssets = SelectedAssets;

		const FUIAction Action_CreateRenderObjectsFromStaticMeshes(
			FExecuteAction::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::ExecuteSelectedContentFunctor, StaticCastSharedPtr<FContentBrowserSelectedAssetExtensionBase>(RenderObjectCreatorFunctor)));
		
		MenuBuilder.AddMenuEntry(
			LOCTEXT("CB_Extension_CreateSoftRenderer", "Create RenderObject"),
			LOCTEXT("CB_Extension_CreateSoftRenderer_Tooltip", "Create RenderObject from selected static meshes"),
			FSlateIcon(FEditorStyle::Get().GetStyleSetName(), "AssetActions.CreateSoftRenderer"),
			Action_CreateRenderObjectsFromStaticMeshes,
			NAME_None,
			EUserInterfaceActionType::Button);
	}

	static void PopulateRenderTexture2DActionsMenu(FMenuBuilder& MenuBuilder, TArray<FAssetData> SelectedAssets)
	{
		const TSharedPtr<FCreateRenderTexture2DFromTextureExtension> RenderObjectCreatorFunctor = MakeShareable(new FCreateRenderTexture2DFromTextureExtension());
		RenderObjectCreatorFunctor->SelectedAssets = SelectedAssets;

		const FUIAction Action_CreateRenderTexture2DFromTextures(
			FExecuteAction::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::ExecuteSelectedContentFunctor, StaticCastSharedPtr<FContentBrowserSelectedAssetExtensionBase>(RenderObjectCreatorFunctor)));
		
		MenuBuilder.AddMenuEntry(
			LOCTEXT("CB_Extension_Texture_CreateRenderTexture2D", "Create RenderTexture2D"),
			LOCTEXT("CB_Extension_Texture_CreateRenderTexture2D_Tooltip", "Create RenderTexture2D from selected textures"),
			FSlateIcon(FEditorStyle::Get().GetStyleSetName(), "AssetActions.CreateSoftRenderer"),
			Action_CreateRenderTexture2DFromTextures,
			NAME_None,
			EUserInterfaceActionType::Button);
	}

	static TSharedRef<FExtender> OnExtendContentBrowserAssetSelectionMenu(const TArray<FAssetData>& SelectedAssets)
	{
		TSharedRef<FExtender> Extender(new FExtender());

		// Run thru the assets to determine if any meet our criteria
		bool bAnyStaticMeshes = false;
		bool bAnyTextures = false;
		for (auto AssetIt = SelectedAssets.CreateConstIterator(); AssetIt; ++AssetIt)
		{
			const FAssetData& Asset = *AssetIt;
			bAnyStaticMeshes = bAnyStaticMeshes || (Asset.AssetClass == UStaticMesh::StaticClass()->GetFName());
			bAnyTextures = bAnyTextures || (Asset.AssetClass == UTexture2D::StaticClass()->GetFName());
		}

		if (bAnyStaticMeshes || bAnyTextures)
		{
			// Add the sprite actions sub-menu extender
			Extender->AddMenuExtension(
				"GetAssetActions",
				EExtensionHook::After,
				nullptr,
				FMenuExtensionDelegate::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::CreateSoftRendererActionsSubMenu, SelectedAssets, bAnyStaticMeshes, bAnyTextures));
		}
		
		return Extender;
	}

	static TArray<FContentBrowserMenuExtender_SelectedAssets>& GetExtenderDelegates()
	{
		FContentBrowserModule& ContentBrowserModule = FModuleManager::LoadModuleChecked<FContentBrowserModule>(TEXT("ContentBrowser"));
		return ContentBrowserModule.GetAllAssetViewContextMenuExtenders();
	}
};

//////////////////////////////////////////////////////////////////////////
// FPaperContentBrowserExtensions

void FSoftRendererContentBrowserExtensions::InstallHooks()
{
	ContentBrowserExtenderDelegate = FContentBrowserMenuExtender_SelectedAssets::CreateStatic(&FSoftRendererContentBrowserExtensions_Impl::OnExtendContentBrowserAssetSelectionMenu);

	TArray<FContentBrowserMenuExtender_SelectedAssets>& CBMenuExtenderDelegates = FSoftRendererContentBrowserExtensions_Impl::GetExtenderDelegates();
	CBMenuExtenderDelegates.Add(ContentBrowserExtenderDelegate);
	ContentBrowserExtenderDelegateHandle = CBMenuExtenderDelegates.Last().GetHandle();
}

void FSoftRendererContentBrowserExtensions::RemoveHooks()
{
	TArray<FContentBrowserMenuExtender_SelectedAssets>& CBMenuExtenderDelegates = FSoftRendererContentBrowserExtensions_Impl::GetExtenderDelegates();
	CBMenuExtenderDelegates.RemoveAll([](const FContentBrowserMenuExtender_SelectedAssets& Delegate){ return Delegate.GetHandle() == ContentBrowserExtenderDelegateHandle; });
}

//////////////////////////////////////////////////////////////////////////

#undef LOCTEXT_NAMESPACE